"""
@Time: 2020/12/17 下午 4:53
@Author: jinzhuan
@File: bert_trigger.py
@Desc: 
"""
from cognlp import *
import torch
import torch.nn as nn
import torch.optim as optim
from cognlp.core.metrics import EventMetric
from cognlp.core.trainer import Trainer
from cognlp.io.loader.ee.ace2005 import ACE2005Loader
from cognlp.io.processor.ee.ace2005_cn import ACE2005CNProcessor
from cognlp.utils.util import get_samples_weight

torch.cuda.set_device(3)
device = torch.device('cuda')

loader = ACE2005Loader()
train_data, dev_data, test_data = loader.load_all('../../../cognlp/data/ee/ace2005-cn/data')
processor = ACE2005CNProcessor(trigger_path='../../../cognlp/data/ee/ace2005-cn/data/trigger_vocabulary.txt',
                               argument_path='../../../cognlp/data/ee/ace2005-cn/data/argument_vocabulary.txt',
                               max_length=256)
trigger_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005-cn/data/trigger_vocabulary.txt')
argument_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005-cn/data/argument_vocabulary.txt')

train_datable = processor.process(train_data)
train_dataset = DataTableSet(train_datable, to_device=False)
train_samples_weight = get_samples_weight(train_datable, trigger_vocabulary)
